from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
from langgraph.graph import MessageGraph

# MessageGraph is deprecated in LangGraph v1.0.0, to be removed in v2.0.0.
# Please use StateGraph with a `messages` key instead.
builder = MessageGraph()

# 向消息图中添加一个节点，节点名称为"chatbot"，该节点执行的函数接收状态并返回一个包含AI消息的列表
# AI消息的内容为"Hello!"，并包含一个名为"search"的工具调用，该工具调用的ID为"123"，参数为"query"为"X"
builder.add_node("chatbot", lambda state: [
    AIMessage(
        content='Hello!',
        tool_calls=[{'name': 'search', 'id': "123", "args": {"query": "X"}}]
    )
])

# 向消息图中添加另一个节点，节点名称为"search"，该节点执行的函数接收状态并返回一个包含工具调用的列表
# 工具消息内容为"Searching for X..."，工具调用的ID为"123"
builder.add_node("search", lambda state: [ToolMessage(content='Searching for X...', tool_call_id="123")])

# 将消息图的入口点设置为"chatbot"
builder.set_entry_point("chatbot")

# 在"chatbot"和"search"节点之间添加一条边
builder.add_edge("chatbot", "search")

# 将"search"节点设置为消息图的结束点
builder.set_finish_point("search")

graph = builder.compile()

print(graph.invoke([HumanMessage(content='Hi there. Can you search for X?')]))
